import pyarrow as pa
import pyarrow.fs
import unittest

from pyarrow import csv


class TestAccessCSV(unittest.TestCase):

    def setUp(self) -> None:
        self.file = 'test.csv'
        self.names = ['name', 'content']
        self.data = [
            pa.array(['name']),
            pa.array(['content'])
        ]
        self.batch = pa.record_batch(self.data, self.names)
        self.local = pyarrow.fs.LocalFileSystem()

    def test_write(self):
        with open(self.file, 'wb') as f:
            writer = csv.CSVWriter(f, self.batch.schema)
            for i in range(50000):
                writer.write(self.batch)

    def test_read(self):
        with open(self.file, 'rb') as f:
            reader = csv.read_csv(f)
            for line in reader:
                print(line)
